import random

from DeterministicMotifFinding import DeterministicMotifFinding
from myseq import MySeq

IUPAC_TABLE = {
        'A': 'A',
        'C': 'C',
        'G': 'G',
        'T': 'T',
        'R': 'AG',
        'Y': 'CT',
        'S': 'CG',
        'W': 'AT',
        'K': 'GT',
        'M': 'AC',
        'B': 'CGT',
        'D': 'AGT',
        'H': 'ACT',
        'V': 'ACG',
        'N': 'ACGT'
    }

class DeterministicMotifFinding_ex(DeterministicMotifFinding):

    def visualize_consensus_motif(self,solution):
        res = ''
        mat = self.create_motif_from_indexes(solution)
        for j in range(self.motif_size):
            max_v,max_char = mat[0][j],self.alphabet[0]
            for i,char in enumerate(self.alphabet[1:]):
                if mat[i+1][j] > max_v:
                    max_v = mat[i+1][j]
                    max_char = char
            res +=  max_char
        return res

    def visualize_consensus_motif_IUPAC(self,solution):
        res = ''
        IUPAC_TABLE_inverted = {v: k for k, v in IUPAC_TABLE.items()}
        mat = self.create_motif_from_indexes(solution)
        for j in range(self.motif_size):
            chars = ''
            for i,char in enumerate(self.alphabet):
                if mat[i][j] != 0:
                    chars += char
            res += IUPAC_TABLE_inverted[chars]
        return res


    def heuristic_consensus_wrapper(self,t):
        for i in range(t):
            random.shuffle(self.seqs)
            sol = self.heuristic_consensus()
            print("{:02}:\tMotif:{}\tScore:{}".format(i+1,
                self.visualize_consensus_motif_IUPAC(sol),
                self.score(sol)))

def test_visualize_consensus_motif():
    seq1 = MySeq("ATAGAGCTGA","DNA")
    seq2 = MySeq("ACGTAGATGA","DNA")
    seq3 = MySeq("AAGATAGGGG","DNA")
    mf = DeterministicMotifFinding_ex(3, [seq1,seq2,seq3])    
    sol = mf.branch_and_bound()
    print("Solution: " , sol)
    print("Motif:", mf.visualize_consensus_motif(sol))
    print("Score:" , mf.score(sol))


def test_visualize_consensus_motif_IUPAC():
    seq1 = MySeq("ACGTAAACCGAAACA","DNA")
    seq2 = MySeq("ACGTAAACCGAAACC","DNA")
    seq3 = MySeq("ACGTCGTGTTCCGGG","DNA")
    seq4 = MySeq("ACGTCGTGTTGTTTT","DNA")
    mf = DeterministicMotifFinding_ex(15, [seq1,seq2,seq3,seq4])    
    sol = mf.branch_and_bound()
    print("Motif IUPAC:", mf.visualize_consensus_motif_IUPAC(sol))

def test_heuristic_consensus_wrapper():
    print("case1:")
    seq1 = MySeq("ATAGAGCTGA","DNA")
    seq2 = MySeq("ACGTAGATGA","DNA")
    seq3 = MySeq("AAGATAGGGG","DNA")
    mf = DeterministicMotifFinding_ex(3, [seq1,seq2,seq3])    
    mf.heuristic_consensus_wrapper(4)

    print("\ncase2:")
    import os
    path = "/bioinformatic_algorithm_colg/7.Motif_Discovery/"
    os.chdir(os.getcwd()+path)    
    mf = DeterministicMotifFinding_ex()
    mf.read_file("exampleMotifs.txt","DNA")

    print ("* Branch and Bound:")
    sol = mf.branch_and_bound()
    print("Motif:", mf.visualize_consensus_motif_IUPAC(sol))
    print("Solution: ", sol)
    print ("Score:", mf.score(sol))
    
    print("* Shuffle Heuristic consensus: ")
    mf.heuristic_consensus_wrapper(10)

if __name__ == "__main__":
    print("### ex1 ###")
    test_visualize_consensus_motif()
    print("### ex2 ###")
    test_visualize_consensus_motif_IUPAC()
    print("### ex3 ###")
    test_heuristic_consensus_wrapper()
